import os
import requests
import gradio as gr
import uuid
import shutil

BACKEND = os.getenv("BACKEND_URL", "http://127.0.0.1:8001/chat")
BACKEND_STREAM = os.getenv("BACKEND_STREAM_URL", "http://127.0.0.1:8001/chat_stream")

# 文件上传目录
UPLOAD_DIR = "../input"
if not os.path.exists(UPLOAD_DIR):
    os.makedirs(UPLOAD_DIR)

def handle_file_upload(file):
    """处理文件上传"""
    if file is None:
        return None, "请选择要上传的文件"
    
    try:
        # 获取文件名
        filename = os.path.basename(file.name)
        # 目标路径
        target_path = os.path.join(UPLOAD_DIR, filename)
        # 复制文件
        shutil.copy2(file.name, target_path)
        return target_path, f"文件已上传：{filename}"
    except Exception as e:
        return None, f"文件上传失败：{str(e)}"

def call_backend_with_file(user_message: str, history: list, uploaded_file_path: str = None):
    """调用后端，支持文件路径"""
    import codecs, time, requests
    accumulated = ""
    decoder = codecs.getincrementaldecoder("utf-8")()
    last_yield_time = 0.0
    
    # 构建请求数据
    request_data = {
        "query": user_message, 
        "history": history, 
        "session_id": SESSION_ID
    }
    
    # 如果有上传的文件，添加文件路径
    if uploaded_file_path:
        request_data["attendance_file_path"] = uploaded_file_path
    
    try:
        with requests.post(
            BACKEND_STREAM,
            json=request_data,
            timeout=120,
            stream=True,
        ) as r:
            r.raise_for_status()
            for chunk in r.iter_content(chunk_size=1024):
                if not chunk:
                    continue
                piece = decoder.decode(chunk)
                if not piece:
                    continue
                for seg in piece.splitlines():
                    if seg.startswith("data:"):
                        seg = seg[5:].lstrip()
                        # 如果是完整的JSON数据，直接返回
                        if seg.strip():
                            accumulated = seg
                            yield accumulated
                            last_yield_time = time.time()
                    else:
                        accumulated += seg
                        now = time.time()
                        if ("\n" in seg) or (len(accumulated) % 20 == 0) or (now - last_yield_time > 0.04):
                            yield accumulated
                            last_yield_time = now
            tail = decoder.decode(b"", final=True)
            if tail:
                accumulated += tail
            if accumulated:
                yield accumulated
    except requests.exceptions.ConnectionError as e:
        yield f"连接不上后端：{BACKEND_STREAM}\n详细：{e}"
    except requests.exceptions.Timeout:
        yield "调用后端超时（>120s）"
    except requests.exceptions.HTTPError as e:
        yield f"后端 HTTP 错误：{getattr(e.response, 'status_code', '?')}\n{getattr(e.response, 'text', '')}"
    except Exception as e:
        yield f"调用后端失败：{e}"

def call_backend(user_message: str, history: list):
    """兼容原有的调用方式"""
    return call_backend_with_file(user_message, history, None)

def create_custom_interface():
    """创建自定义界面，包含文件上传功能"""
    with gr.Blocks(title="劳动关系大模型___收费版") as demo:
        gr.Markdown("# 劳动关系大模型___收费版")
        gr.Markdown(f"流式后端：{BACKEND_STREAM}")
        
        with gr.Row():
            with gr.Column(scale=3):
                # 聊天界面
                chatbot = gr.Chatbot(height=500, type="messages")
                msg = gr.Textbox(
                    label="输入消息",
                    placeholder="请输入您的问题，或上传文件后询问分析...",
                    lines=2
                )
                
                with gr.Row():
                    submit_btn = gr.Button("发送", variant="primary")
                    clear_btn = gr.Button("清空对话")
            
            with gr.Column(scale=1):
                # 文件上传区域
                gr.Markdown("### 📁 文件上传")
                file_upload = gr.File(
                    label="上传文件",
                    file_types=[".pdf", ".docx", ".jpg", ".jpeg", ".png"],
                    type="filepath"
                )
                upload_status = gr.Textbox(
                    label="上传状态",
                    interactive=False,
                    lines=2
                )
                
                # 劳动赔偿计算器区域
                gr.Markdown("### ⚖️ 劳动赔偿计算器")
                
                # 计算器显示控制
                calculator_visible = gr.State(False)
                
                with gr.Group(visible=False) as calculator_group:
                    compensation_type = gr.Dropdown(
                        label="赔偿类型",
                        choices=[
                            "经济补偿金（N）",
                            "赔偿金（2N）",
                            "双倍工资",
                            "拖欠工资赔偿",
                            "工伤医疗赔偿",
                            "试用期违法延长赔偿",
                            "未休年假工资",
                            "生育津贴/产假工资差额",
                            "加班费",
                            "违法调岗降薪赔偿"
                        ],
                        value="经济补偿金（N）"
                    )
                    
                    city_select = gr.Dropdown(
                        label="城市",
                        choices=[
                            "北京", "上海", "广州", "深圳", "杭州", "南京", "苏州", "成都", "武汉", "西安", "天津", "重庆", "青岛", "大连", "宁波", "厦门", "长沙"
                        ],
                        value="北京"
                    )
                    
                    monthly_salary = gr.Number(
                        label="月平均工资（元）",
                        value=8000,
                        minimum=0
                    )
                    
                    work_years = gr.Number(
                        label="工作年限（年）",
                        value=3,
                        minimum=0,
                        step=0.1
                    )
                    
                    # 动态参数字段
                    with gr.Row(visible=False) as maternity_params:
                        maternity_allowance = gr.Number(
                            label="生育津贴（元/月）",
                            value=6000,
                            minimum=0
                        )
                        maternity_months = gr.Number(
                            label="产假月数",
                            value=4,
                            minimum=0,
                            step=0.1
                        )
                    
                    with gr.Row(visible=False) as overtime_params:
                        weekday_hours = gr.Number(
                            label="工作日加班时数",
                            value=20,
                            minimum=0
                        )
                        weekend_hours = gr.Number(
                            label="周末加班时数",
                            value=8,
                            minimum=0
                        )
                        holiday_hours = gr.Number(
                            label="节假日加班时数",
                            value=4,
                            minimum=0
                        )
                    
                    with gr.Row(visible=False) as injury_params:
                        medical_expenses = gr.Number(
                            label="医疗费用（元）",
                            value=10000,
                            minimum=0
                        )
                        disability_level = gr.Number(
                            label="伤残等级（1-10级）",
                            value=7,
                            minimum=1,
                            maximum=10,
                            step=1
                        )
                    
                    with gr.Row(visible=False) as leave_params:
                        unused_days = gr.Number(
                            label="未休年假天数",
                            value=5,
                            minimum=0,
                            step=1
                        )
                    
                    with gr.Row(visible=False) as salary_params:
                        unsigned_months = gr.Number(
                            label="未签合同月数",
                            value=12,
                            minimum=0,
                            step=1
                        )
                    
                    with gr.Row(visible=False) as overdue_params:
                        overdue_amount = gr.Number(
                            label="拖欠工资金额（元）",
                            value=15000,
                            minimum=0
                        )
                    
                    with gr.Row(visible=False) as probation_params:
                        probation_salary = gr.Number(
                            label="试用期工资（元）",
                            value=4000,
                            minimum=0
                        )
                        excess_days = gr.Number(
                            label="超期天数",
                            value=30,
                            minimum=1,
                            maximum=365
                        )
                    
                    with gr.Row(visible=False) as transfer_params:
                        new_salary = gr.Number(
                            label="调岗后工资（元）",
                            value=4000,
                            minimum=0
                        )
                        reduction_months = gr.Number(
                            label="降薪月数",
                            value=6,
                            minimum=1,
                            maximum=60
                        )
                    
                    calculate_btn = gr.Button("计算赔偿", variant="primary")
                    
                    calculation_result = gr.Textbox(
                        label="计算结果",
                        interactive=False,
                        lines=8
                    )
                
                # 示例问题
                gr.Markdown("### 💡 示例问题")
                examples = gr.Examples(
                    examples=[
                        "上海的气温怎么样？",
                        "工资被拖欠了怎么办？",
                        "我想分析一下我的考勤记录",
                        "请帮我审查这份合同",
                        "请分析这份工资单",
                        "我要计算劳动赔偿"
                    ],
                    inputs=msg
                )
        
        # 存储上传的文件路径
        uploaded_file_path = gr.State(None)
        
        def handle_file_upload_ui(file):
            """处理文件上传的UI函数"""
            if file is None:
                return None, "请选择要上传的文件"
            
            try:
                # 获取文件名
                filename = os.path.basename(file)
                # 目标路径
                target_path = os.path.join(UPLOAD_DIR, filename)
                # 复制文件
                shutil.copy2(file, target_path)
                return target_path, f"✅ 文件已上传：{filename}\n路径：{target_path}"
            except Exception as e:
                return None, f"❌ 文件上传失败：{str(e)}"
        
        def respond(message, chat_history, file_path):
            """处理用户消息并返回响应"""
            if not message.strip():
                return "", chat_history, gr.update(visible=False)
            
            # 添加用户消息到历史
            chat_history.append({"role": "user", "content": message})
            
            # 调用后端
            accumulated = ""
            show_calculator = False
            
            for chunk in call_backend_with_file(message, [], file_path):
                accumulated = chunk
                
            # 在最终的accumulated中尝试解析JSON格式的响应
            try:
                import json
                json_data = json.loads(accumulated)
                if isinstance(json_data, dict):
                    accumulated = json_data.get('content', accumulated)
                    if json_data.get('show_calculator', False):
                        show_calculator = True
            except json.JSONDecodeError:
                # 如果不是JSON格式，检查是否包含劳动赔偿关键词
                labor_keywords = ["劳动赔偿", "经济补偿", "赔偿金", "计算赔偿", "补偿计算", "加班费"]
                show_calculator = any(keyword in message for keyword in labor_keywords)
            
            # 如果是劳动赔偿相关请求，显示计算器
            if show_calculator:
                accumulated += "\n\n💡 **提示：右侧已为您展开劳动赔偿计算器，您可以直接填写参数进行计算！**"
            
            # 添加助手回复到历史
            chat_history.append({"role": "assistant", "content": accumulated})
            
            return "", chat_history, gr.update(visible=show_calculator)
        
        def clear_chat():
            """清空对话历史"""
            return [], None, "对话已清空", gr.update(visible=False)
        
        def update_param_visibility(comp_type):
            """根据赔偿类型更新参数字段的可见性"""
            return {
                maternity_params: gr.update(visible=(comp_type == "生育津贴/产假工资差额")),
                overtime_params: gr.update(visible=(comp_type == "加班费")),
                injury_params: gr.update(visible=(comp_type == "工伤医疗赔偿")),
                leave_params: gr.update(visible=(comp_type == "未休年假工资")),
                salary_params: gr.update(visible=(comp_type == "双倍工资")),
                overdue_params: gr.update(visible=(comp_type == "拖欠工资赔偿")),
                probation_params: gr.update(visible=(comp_type == "试用期违法延长赔偿")),
                transfer_params: gr.update(visible=(comp_type == "违法调岗降薪赔偿"))
            }
        
        def calculate_compensation(comp_type, city, salary, years, mat_allowance, mat_months, 
                                 wd_hours, we_hours, hd_hours, med_exp, dis_level, 
                                 un_days, un_months, od_amount, prob_salary, ex_days, 
                                 new_sal, red_months):
            """计算劳动赔偿"""
            try:
                # 导入计算器
                import sys
                import os
                sys.path.append(os.path.dirname(__file__))
                from 劳动赔偿计算器 import LaborCompensationCalculator
                
                calculator = LaborCompensationCalculator()
                
                # 映射赔偿类型
                type_mapping = {
                    "经济补偿金（N）": "economic_compensation",
                    "赔偿金（2N）": "severance_pay",
                    "双倍工资": "double_salary",
                    "拖欠工资赔偿": "overdue_salary",
                    "工伤医疗赔偿": "work_injury",
                    "试用期违法延长赔偿": "probation_extension",
                    "未休年假工资": "annual_leave",
                    "生育津贴/产假工资差额": "maternity_allowance",
                    "加班费": "overtime",
                    "违法调岗降薪赔偿": "illegal_transfer"
                }
                
                compensation_type = type_mapping.get(comp_type, "economic_compensation")
                
                # 调用具体的计算方法
                result = None
                if comp_type == "经济补偿金（N）":
                    result = calculator.calculate_economic_compensation(salary, years, city)
                elif comp_type == "赔偿金（2N）":
                    result = calculator.calculate_severance_pay(salary, years, city)
                elif comp_type == "双倍工资":
                    result = calculator.calculate_double_salary(salary, un_months or 12)
                elif comp_type == "拖欠工资赔偿":
                    result = calculator.calculate_overdue_salary_compensation(od_amount or salary)
                elif comp_type == "工伤医疗赔偿":
                    result = calculator.calculate_work_injury_compensation(med_exp or 10000, dis_level or 7, salary)
                elif comp_type == "试用期违法延长赔偿":
                    result = calculator.calculate_probation_extension_compensation(prob_salary or salary*0.8, salary, ex_days or 30)
                elif comp_type == "未休年假工资":
                    result = calculator.calculate_annual_leave_compensation(salary, un_days or 5)
                elif comp_type == "生育津贴/产假工资差额":
                    result = calculator.calculate_maternity_allowance_difference(salary, mat_allowance or salary*0.8, mat_months or 4)
                elif comp_type == "加班费":
                    result = calculator.calculate_overtime_compensation(salary, wd_hours or 20, we_hours or 8, hd_hours or 4)
                elif comp_type == "违法调岗降薪赔偿":
                    result = calculator.calculate_illegal_transfer_compensation(salary, new_sal or salary*0.8, red_months or 6, years)
                else:
                    result = calculator.calculate_economic_compensation(salary, years, city)
                
                if result:
                    # 格式化结果
                    output = f"🎯 **{comp_type}计算结果**\n\n"
                    output += f"📍 **城市：** {city}\n"
                    output += f"💰 **月平均工资：** {salary:,.2f} 元\n"
                    output += f"⏰ **工作年限：** {years} 年\n\n"
                    output += f"💵 **应得赔偿：** {result['amount']:,.2f} 元\n\n"
                    output += f"📋 **计算公式：**\n{result['formula']}\n\n"
                    output += f"⚖️ **法律依据：**\n{result['legal_basis']}\n\n"
                    if result.get('notes'):
                        output += f"📝 **说明：**\n"
                        for note in result['notes']:
                            output += f"• {note}\n"
                    
                    return output
                else:
                    return f"❌ 计算失败：{result.get('message', '未知错误')}"
                    
            except Exception as e:
                return f"❌ 计算出错：{str(e)}"
        
        # 事件绑定
        file_upload.change(
            handle_file_upload_ui,
            inputs=[file_upload],
            outputs=[uploaded_file_path, upload_status]
        )
        
        submit_btn.click(
            respond,
            inputs=[msg, chatbot, uploaded_file_path],
            outputs=[msg, chatbot, calculator_group]
        )
        
        msg.submit(
            respond,
            inputs=[msg, chatbot, uploaded_file_path],
            outputs=[msg, chatbot, calculator_group]
        )
        
        clear_btn.click(
            clear_chat,
            outputs=[chatbot, uploaded_file_path, upload_status, calculator_group]
        )
        
        # 计算器事件绑定
        compensation_type.change(
            update_param_visibility,
            inputs=[compensation_type],
            outputs=[maternity_params, overtime_params, injury_params, leave_params, salary_params, overdue_params, probation_params, transfer_params]
        )
        
        calculate_btn.click(
            calculate_compensation,
            inputs=[
                compensation_type, city_select, monthly_salary, work_years,
                maternity_allowance, maternity_months,
                weekday_hours, weekend_hours, holiday_hours,
                medical_expenses, disability_level,
                unused_days, unsigned_months, overdue_amount,
                probation_salary, excess_days, new_salary, reduction_months
            ],
            outputs=[calculation_result]
        )
    
    return demo

if __name__ == "__main__":
    os.environ["NO_PROXY"] = os.environ.get("NO_PROXY", "") + ",127.0.0.1,localhost"
    os.environ["no_proxy"] = os.environ["NO_PROXY"]

    SESSION_ID = str(uuid.uuid4())
    print(f"[INFO] BACKEND_STREAM = {BACKEND_STREAM}")
    # 可打印一下版本以便排查: print("gradio version:", gr.__version__)

    demo = create_custom_interface()
    demo.launch(server_name="127.0.0.1", server_port=7860)
